// Package db : 数据库连接模块
// 定义数据库的连接和基础操作
// 封装数据库的基础操作，如增删改查等，方便后续开发
// 通过连接池管理相同IP地址和端口号的数据库连接
package db

import (
	"database/sql"
	"fmt"
	"strings"
	"sync"
	"time"

	"distributed_db/utils/transaction"

	_ "github.com/go-sql-driver/mysql"
)

// DBConnection : 数据库连接对象结构体
type DBConnection struct {
	host     string
	port     string
	user     string
	password string
	dbName   string
	db       *sql.DB
	refCount int
	mu       sync.Mutex
}

// Region : 表示一个 Region，即 contributed_db 数据库中的一个表
// RegionName 表示 Region 名称（即表名）
// RecordNum 表示 Region 的大小（即表中的记录数量）
type Region struct {
	RegionName string
	RecordNum  int
}

// connectionPool : 全局连接池，用于管理相同IP地址和端口号的数据库连接
var connectionPool = sync.Map{}

// NewDBConnection : 创建或复用数据库连接对象
func NewDBConnection(host, port, user, password, dbName string) *DBConnection {
	key := host + ":" + port + ":" + dbName
	if conn, exists := connectionPool.Load(key); exists {
		conn.(*DBConnection).mu.Lock()
		conn.(*DBConnection).refCount++
		conn.(*DBConnection).mu.Unlock()
		return conn.(*DBConnection)
	}
	conn := &DBConnection{
		host:     host,
		port:     port,
		user:     user,
		password: password,
		dbName:   dbName,
	}
	connectionPool.Store(key, conn)
	return conn
}

// GetConnection : 获取数据库连接
func (db *DBConnection) GetConnection() (*sql.DB, error) {
	db.mu.Lock()
	defer db.mu.Unlock()

	if db.db == nil {
		// 连接到指定的 Mysql 服务
		dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/", db.user, db.password, db.host, db.port)
		var err error
		db.db, err = sql.Open("mysql", dsn)
		if err != nil {
			return nil, fmt.Errorf("failed to open database connection: %v", err)
		}

		// 确保连接到指定的数据库，如果不存在则创建
		_, err = db.db.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", db.dbName))
		if err != nil {
			return nil, fmt.Errorf("failed to create database: %v", err)
		}

		// 数据库连接池相关设置
		db.db.SetMaxOpenConns(20)
		db.db.SetMaxIdleConns(20)
		db.db.SetConnMaxLifetime(time.Minute * 5)

		// 使用该数据库
		_, err = db.db.Exec(fmt.Sprintf("USE %s", db.dbName))
		if err != nil {
			return nil, fmt.Errorf("failed to switch to database: %v", err)
		}
	}
	return db.db, nil
}

// ReleaseConnection : 释放数据库连接
func (db *DBConnection) ReleaseConnection() {
	db.mu.Lock()
	defer db.mu.Unlock()

	db.refCount--
	if db.refCount == 0 {
		db.closeConnection()
	}
}

// closeConnection : 关闭数据库连接
func (db *DBConnection) closeConnection() {
	if db.db != nil {
		db.db.Close()
		db.db = nil
	}
}

// tableExists : 检查表是否存在
func (db *DBConnection) tableExists(conn *sql.DB, table string) (bool, error) {
	var count int
	query := fmt.Sprintf("SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = '%s' AND table_name = '%s'", db.dbName, table)
	err := conn.QueryRow(query).Scan(&count)
	if err != nil {
		return false, err
	}
	return count > 0, nil
}

// createTable : 创建表
// 参数：
// * table: 表名
// * columns: 字段列表
// * primaryKey: 主键字段列表
// * indexes: 索引列表（新增参数）
// 返回值：
// * error：错误信息
func (db *DBConnection) createTable(conn *sql.DB, table string, columns []string, primaryKey []string, indexes []string) error {
	columnDefinitions := ""
	for _, col := range columns {
		if columnDefinitions != "" {
			columnDefinitions += ", "
		}
		columnDefinitions += fmt.Sprintf("%s VARCHAR(255)", col)
	}

	// 如果主键字段不为空，则添加主键定义
	if len(primaryKey) > 0 {
		columnDefinitions += ", PRIMARY KEY (" + strings.Join(primaryKey, ", ") + ")"
	}

	// 切换到对应数据库
	_, err := conn.Exec(fmt.Sprintf("USE %s", db.dbName))
	if err != nil {
		return err
	}

	// 运行建表语句
	query := fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s)", table, columnDefinitions)
	_, err = conn.Exec(query)
	if err != nil {
		return fmt.Errorf("failed to create table: %v", err)
	}

	// 创建索引
	if len(indexes) > 0 {
		err = db.CreateIndex(table, indexes)
		if err != nil {
			return fmt.Errorf("failed to create indexes: %v", err)
		}
	}

	return nil
}

// insertData : 插入数据
// 参数：
// * table: 表名
// * columns: 字段列表
// * values: 字段值列表（二维数组，支持批量插入）
// 返回值：
// * error：错误信息
func (db *DBConnection) insertData(conn *sql.DB, tx *sql.Tx, table string, columns []string, values [][]interface{}) (int, error) {
	if len(values) == 0 {
		return 0, nil
	}

	// 构造字段名部分
	columnsStr := ""
	for i, col := range columns {
		if i > 0 {
			columnsStr += ", "
		}
		columnsStr += col
	}

	// 构造占位符部分
	placeholder := "(" + "?" + strings.Repeat(", ?", len(columns)-1) + ")"

	// 构造批量插入的值部分
	valuePlaceholders := ""
	var args []interface{}
	for i, row := range values {
		if i > 0 {
			valuePlaceholders += ", "
		}
		valuePlaceholders += placeholder
		args = append(args, row...)
	}

	query := fmt.Sprintf("INSERT INTO %s (%s) VALUES %s", table, columnsStr, valuePlaceholders)

	//fmt.Printf("query:", query)

	var err error

	if tx != nil {
		_, err = tx.Exec(fmt.Sprintf("USE %s", db.dbName))
		_, err = tx.Exec(query, args...)
	} else {
		_, err = conn.Exec(fmt.Sprintf("USE %s", db.dbName))
		_, err = conn.Exec(query, args...)
	}
	if err != nil {
		return 0, err
	}

	return len(values), nil
}

// BeginTx 开启一个新的事务
func (db *DBConnection) BeginTx() (*sql.Tx, error) {
	conn, err := db.GetConnection()
	if err != nil {
		return nil, err
	}
	return conn.Begin()
}

// InsertRegion : 创建在指定数据库中创建 Region（Table）
// 参数：
// * table: 操作的表名
// * columns: 表中所有的字段名
// * primaryKeys: 表中字段的主键
// * indexes: 表中建立索引的字段
// * values: 表中所有字段对应的值（二维数组，支持批量插入）
// 返回值
// * error：错误信息
func (db *DBConnection) InsertRegion(tx *transaction.Transaction, table string, columns []string, primaryKeys []string, indexes []string, values [][]interface{}) (int, error) {
	conn, err := db.GetConnection()
	if err != nil {
		fmt.Printf("\033[31mInsertRegion: failed to get connection: %v\033[0m\n", err)
		return 0, err
	}
	defer db.ReleaseConnection()

	var dbTx *sql.Tx
	if tx != nil {
		tx.TxMapLock.Lock()
		// 确保 TxMap 已初始化
		if tx.TxMap == nil {
			tx.TxMap = make(map[string]*sql.Tx)
		}
		dbTx = tx.TxMap[db.host+":"+db.port+":"+db.dbName]
		if dbTx == nil {
			dbTx, err = db.BeginTx()
			if err != nil {
				tx.TxMapLock.Unlock()
				return 0, fmt.Errorf("failed to begin transaction: %v", err)
			}
			tx.TxMap[db.host+":"+db.port+":"+db.dbName] = dbTx
		}
		tx.TxMapLock.Unlock()
	} else {
		_, err := conn.Exec(fmt.Sprintf("USE %s", db.dbName))
		if err != nil {
			return 0, fmt.Errorf("failed to switch to database: %v", err)
		}
	}

	// 检查表是否存在，如果不存在则创建表
	exists, err := db.tableExists(conn, table)
	if err != nil {
		fmt.Printf("\033[31mInsertRegion: failed to check if table exists: %v \033[0m\n", err)
		return 0, fmt.Errorf("failed to check if table exists: %v", err)
	}
	if !exists {
		if err := db.createTable(conn, table, columns, primaryKeys, indexes); err != nil {
			fmt.Printf("\033[31mInsertRegion: failed to create table: %v \033[0m\n", err)
			return 0, err
		}
	}
	// 插入数据时使用表锁
	if tx != nil {
		_, err = dbTx.Exec(fmt.Sprintf("LOCK TABLES %s WRITE", table))
		defer dbTx.Exec(fmt.Sprintf("UNLOCK TABLES"))
	} else {
		_, err = conn.Exec(fmt.Sprintf("LOCK TABLES %s WRITE", table))
		defer conn.Exec(fmt.Sprintf("UNLOCK TABLES"))
	}

	if err != nil {
		fmt.Printf("\033[31mInsertRegion: failed to lock table for write: %v\033[0m\n", err)
		return 0, fmt.Errorf("failed to lock table for write: %v", err)
	}

	inserted, err := db.insertData(conn, dbTx, table, columns, values)
	if err != nil {
		fmt.Printf("\033[31mInsertRegion: failed to insert data: %v\033[0m\n", err)
		return 0, err
	}

	return inserted, nil
}

// QueryRegion : 根据传入的条件，查询指定数据库中指定表的数据
func (db *DBConnection) QueryRegion(tx *transaction.Transaction, table, queryRow, condition string) ([]map[string]interface{}, error) {
	conn, err := db.GetConnection()
	if err != nil {
		return nil, err
	}
	defer db.ReleaseConnection()

	var dbTx *sql.Tx
	if tx != nil {
		tx.TxMapLock.Lock()
		dbTx, _ := tx.TxMap[db.host+":"+db.port+":"+db.dbName]
		if dbTx == nil {
			dbTx, err = db.BeginTx()
			if err != nil {
				tx.TxMapLock.Unlock()
				return nil, fmt.Errorf("failed to begin transaction: %v", err)
			}
			tx.TxMap[db.host+":"+db.port+":"+db.dbName] = dbTx
		}
		tx.TxMapLock.Unlock()
	} else {
		_, err := conn.Exec(fmt.Sprintf("USE %s", db.dbName))
		if err != nil {
			return nil, fmt.Errorf("failed to switch to database: %v", err)
		}
	}

	// 检查表是否存在
	exists, err := db.tableExists(conn, table)
	if err != nil {
		return nil, fmt.Errorf("failed to check if table exists: %v", err)
	}
	if !exists {
		return nil, fmt.Errorf("table %s does not exist", table)
	}

	// 查询数据时使用共享锁
	_, err = conn.Exec(fmt.Sprintf("LOCK TABLES %s READ", table))
	if err != nil {
		return nil, fmt.Errorf("failed to lock table for read: %v", err)
	}
	defer conn.Exec("UNLOCK TABLES")

	var query string
	// 查询表中的所有数据
	if condition == "" {
		query = fmt.Sprintf("SELECT %s FROM %s", queryRow, table)
	} else {
		query = fmt.Sprintf("SELECT %s FROM %s WHERE %s", queryRow, table, condition)
	}

	var rows *sql.Rows
	if dbTx != nil {
		_, err = dbTx.Exec(fmt.Sprintf("USE %s", db.dbName))
		rows, err = dbTx.Query(query)
	} else {
		_, err = conn.Exec(fmt.Sprintf("USE %s", db.dbName))
		rows, err = conn.Query(query)
	}
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	var results []map[string]interface{}
	for rows.Next() {
		row := make([]interface{}, len(columns))
		rowPointers := make([]interface{}, len(columns))
		for i := range columns {
			rowPointers[i] = &row[i]
		}
		if err := rows.Scan(rowPointers...); err != nil {
			return nil, err
		}
		result := make(map[string]interface{})
		for i, col := range columns {
			result[col] = row[i]
		}
		results = append(results, result)
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	//fmt.Printf("db results:", results)

	return results, nil
}

// DeleteRegion : 根据传入的条件，删除指定数据库中指定表的数据
func (db *DBConnection) DeleteRegion(tx *transaction.Transaction, table, condition string) (int, error) {
	conn, err := db.GetConnection()
	if err != nil {
		return 0, err
	}
	defer db.ReleaseConnection()

	var dbTx *sql.Tx
	if tx != nil {
		tx.TxMapLock.Lock()
		dbTx, _ = tx.TxMap[db.host+":"+db.port+":"+db.dbName]
		tx.TxMapLock.Unlock()
	}

	// 检查表是否存在
	exists, err := db.tableExists(conn, table)
	if err != nil {
		return 0, fmt.Errorf("failed to check if table exists: %v", err)
	}
	if !exists {
		return 0, fmt.Errorf("table %s does not exist", table)
	}

	// 删除数据
	query := fmt.Sprintf("DELETE FROM %s WHERE %s", table, condition)
	var result sql.Result
	var errDelete error
	if dbTx != nil {
		_, errDelete = dbTx.Exec(fmt.Sprintf("USE %s", db.dbName))
		result, errDelete = dbTx.Exec(query)
	} else {
		_, errDelete = conn.Exec(fmt.Sprintf("USE %s", db.dbName))
		result, errDelete = conn.Exec(query)
	}
	if errDelete != nil {
		return 0, fmt.Errorf("failed to delete data: %v", errDelete)
	}

	// 获取删除的行数
	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("failed to get rows affected: %v", err)
	}

	return int(rowsAffected), nil
}

// UpdateRegion : 根据传入的条件，更新指定数据库中指定表的数据
func (db *DBConnection) UpdateRegion(tx *transaction.Transaction, table, clause, condition string) error {
	conn, err := db.GetConnection()
	if err != nil {
		return err
	}
	defer db.ReleaseConnection()

	var dbTx *sql.Tx
	if tx != nil {
		tx.TxMapLock.Lock()
		dbTx, _ = tx.TxMap[db.host+":"+db.port+":"+db.dbName]
		tx.TxMapLock.Unlock()
	}

	// 检查表是否存在
	exists, err := db.tableExists(conn, table)
	if err != nil {
		return fmt.Errorf("failed to check if table exists: %v", err)
	}
	if !exists {
		return fmt.Errorf("table %s does not exist", table)
	}

	// 检查是否存在满足条件的数据
	var count int
	query := fmt.Sprintf("SELECT COUNT(*) FROM %s WHERE %s", table, condition)
	err = conn.QueryRow(query).Scan(&count)
	if err != nil {
		return fmt.Errorf("failed to check if data exists: %v", err)
	}
	if count == 0 {
		// 如果不存在满足条件的数据，则不进行任何更新操作
		return nil
	}

	query = fmt.Sprintf("UPDATE %s SET %s WHERE %s", table, clause, condition)

	//fmt.Printf("query is: ", query)

	// 执行更新操作
	var errUpdate error
	if dbTx != nil {
		_, errUpdate = dbTx.Exec(fmt.Sprintf("USE %s", db.dbName))
		_, errUpdate = dbTx.Exec(query)
	} else {
		_, errUpdate = conn.Exec(fmt.Sprintf("USE %s", db.dbName))
		_, errUpdate = conn.Exec(query)
	}
	if errUpdate != nil {
		return fmt.Errorf("failed to update data: %v", errUpdate)
	}

	//fmt.Printf("Updated region table: ", table)

	return nil
}

// GetRegionData : 获取Region（表单）中的所有数据
// 参数：
// * table: 操作的表名
// 返回值
// * []string：字段列表
// * [][]interface{}：字段值列表
// * error：错误信息
// 操作逻辑：
//  1. 获取数据库连接
//  2. 根据传入的表名，对该表执行全表全列数据的查询操作
//     如果该表不存在，则返回相应的报错信息
//  3. 将数据库连接返回到连接池中，方便下次使用
func (db *DBConnection) GetRegionData(table string) ([]string, [][]interface{}, error) {
	conn, err := db.GetConnection()
	if err != nil {
		return nil, nil, err
	}
	defer db.ReleaseConnection()

	// 检查表是否存在
	exists, err := db.tableExists(conn, table)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to check if table exists: %v", err)
	}
	if !exists {
		return nil, nil, fmt.Errorf("table %s does not exist", table)
	}

	// 查询数据时使用共享锁
	_, err = conn.Exec(fmt.Sprintf("LOCK TABLES %s READ", table))
	if err != nil {
		return nil, nil, fmt.Errorf("failed to lock table for read: %v", err)
	}
	defer conn.Exec("UNLOCK TABLES")

	// 查询表中的所有数据
	query := fmt.Sprintf("SELECT * FROM %s", table)
	rows, err := conn.Query(query)
	if err != nil {
		return nil, nil, err
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil {
		return nil, nil, err
	}

	var results [][]interface{}
	for rows.Next() {
		row := make([]interface{}, len(columns))
		rowPointers := make([]interface{}, len(columns))
		for i := range columns {
			rowPointers[i] = &row[i]
		}
		if err := rows.Scan(rowPointers...); err != nil {
			return nil, nil, err
		}
		results = append(results, row)
	}

	if err := rows.Err(); err != nil {
		return nil, nil, err
	}

	return columns, results, nil
}

// GetAllRegions : 获取指定 Region Server 中的所有 Region（即指定数据库中的所有表）
// 参数：
// * 无
// 返回值
// * []Region：Region 列表
// * error：错误信息
func (db *DBConnection) GetAllRegions() ([]Region, error) {
	conn, err := db.GetConnection()
	if err != nil {
		return nil, err
	}
	defer db.ReleaseConnection()

	rows, err := conn.Query(fmt.Sprintf("SHOW TABLES FROM %s", db.dbName))
	_, err = conn.Exec("USE " + db.dbName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var regions []Region
	for rows.Next() {
		var region Region
		if err := rows.Scan(&region.RegionName); err != nil {
			return nil, err
		}

		// 获取表的记录数量
		var recordCount int
		query := fmt.Sprintf("SELECT COUNT(*) FROM %s", region.RegionName)
		err = conn.QueryRow(query).Scan(&recordCount)
		if err != nil {
			return nil, err
		}
		region.RecordNum = recordCount
		regions = append(regions, region)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	return regions, nil
}

// DropRegion : 在指定的 RegionServer 上删除指定的 Region
func (db *DBConnection) DropRegion(regionName string) error {
	conn, err := db.GetConnection()
	if err != nil {
		return err
	}
	defer db.ReleaseConnection()

	_, err = conn.Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s", regionName))
	if err != nil {
		return err
	}

	return nil
}

// DropAllRegions : 在指定的 RegionServer 上删除所有 Region（Table）
func (db *DBConnection) DropAllRegions() error {
	conn, err := db.GetConnection()
	if err != nil {
		return err
	}
	defer db.ReleaseConnection()

	rows, err := conn.Query(fmt.Sprintf("SHOW TABLES FROM %s", db.dbName))
	if err != nil {
		return err
	}
	defer func(rows *sql.Rows) {
		err := rows.Close()
		if err != nil {
			return
		}
	}(rows)

	_, err = db.db.Exec(fmt.Sprintf("USE %s", db.dbName))
	if err != nil {
		return err
	}

	var tableName string
	for rows.Next() {
		if err := rows.Scan(&tableName); err != nil {
			return err
		}
		_, err := conn.Exec(fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName))
		if err != nil {
			return err
		}
	}
	if err := rows.Err(); err != nil {
		return err
	}

	return nil
}

// CreateIndex : 在指定表的列上创建索引
// 参数：
// * table: 表名
// * indexes: 索引列表，每个索引是一个字符串，表示索引名称和对应的列
// 返回值：
// * error：错误信息
func (db *DBConnection) CreateIndex(table string, indexes []string) error {
	for _, index := range indexes {
		indexParts := strings.Split(index, ":")
		if len(indexParts) != 2 {
			return fmt.Errorf("invalid index format: %s", index)
		}
		indexName := indexParts[0]
		indexColumn := indexParts[1]

		query := fmt.Sprintf("CREATE INDEX %s ON %s (%s)", indexName, table, indexColumn)
		_, err := db.db.Exec(query)
		if err != nil {
			return fmt.Errorf("failed to create index %s on table %s: %v", indexName, table, err)
		}
	}
	return nil
}

// ReleaseAllConnections : 关闭连接池中的所有连接
func ReleaseAllConnections() {
	connectionPool.Range(func(key, value interface{}) bool {
		dbConn := value.(*DBConnection)
		dbConn.closeConnection()
		return true
	})
}
